import numpy as np
import math
import matplotlib.pyplot as plt

from numpy.lib.npyio import savez_compressed

class Config():
    def __init__(self):
        self.max_speed = 1.0
        self.min_speed = 0
        self.max_yawrate = 40.0 * math.pi / 180.0  # [rad/s]
        self.accel = 0.2
        self.max_dyawrate = 40.0 * math.pi/180.0
        self.v_reso = 0.01
        self.yawrate_reso = 0.1 * math.pi / 180.0
        self.dt = 0.1
        self.predict_time =3
        self.to_goal_cost_gain = 10
        self.speed_cost_gain = 1.0
        self.obstacle_cost_gain =1.0
        self.robot_radius = 0.1


def predict_trajectory(x_init,v,y,config):
    x = np.array(x_init)
    traj = np.array(x)
    time = 0
    while time <=config.predict_time:
        x = motion(x,[v,y],config.dt)
        traj = np.vstack((traj,x))
        time +=config.dt
    return traj

def calc_to_goal_cost(traj, goal, config):
    dx = goal[0] -traj[-1][0]
    dy = goal[1] -traj[-1][1]
    # cost 越小越好
    cost = math.sqrt(dx**2 + dy**2)*100
    return cost


def calc_to_angel_cost(traj, goal, config):
    dx = goal[0] -traj[-1][0]
    dy = goal[1] -traj[-1][1]
    # cost 越小越好
    cost = math.sqrt(dx**2 + dy**2)*100
    return cost

def calc_obstacle_cost(traj,ob,config):
    skip_n = 2
    minr = float("inf")
    for ii in range(0,len(traj[:,1]),skip_n):
        for i in range(len(ob[:,0])):
            ox=ob[i,0]
            oy=ob[i,1]
            dx = traj[ii,0]-ox
            dy = traj[ii,1]-oy
            r = math.sqrt(dx**2+dy**2)
            if r < config.robot_radius:
                return float("Inf")
            if minr >= r:
                minr = r
            return 1.0/minr



def calc_final_input(x,u,dw,config,goal,ob):
    x_init = x[:]
    min_cost = float("inf")
    best_u = [0.0,0.0]
    best_traj = np.array([x])
    ## 对生成的轨迹进行评估
    for v in np.arange(dw[0],dw[1],config.v_reso):
        for y in np.arange(dw[2],dw[3],config.yawrate_reso):
            
            traj = predict_trajectory(x_init,v,y,config)
            to_goal_cost = config.to_goal_cost_gain*calc_to_goal_cost(traj,goal,config)
            speed_cost = config.speed_cost_gain*(config.max_speed - traj[-1][3])
            ob_cost = config.obstacle_cost_gain*calc_obstacle_cost(traj,ob,config)
            final_cost = to_goal_cost*1+speed_cost+ob_cost

            if final_cost < min_cost:
                best_u = [v,y]
                best_traj = traj
                min_cost = final_cost
    return best_u,best_traj

## 给出当前可以达到的最大速度
def calc_dynamic_window(x,config):
    speed_x_max = x[3] + config.accel*config.dt
    speed_x_min = x[3] - config.accel*config.dt
    speed_y_max = x[4] + config.max_dyawrate*config.dt
    speed_y_min = x[4] - config.max_dyawrate*config.dt
    Vd = [speed_x_min,speed_x_max,speed_y_min,speed_y_max]
    Vs = [config.min_speed,config.max_speed,-config.max_yawrate,config.max_yawrate]
    
    dw = [max(Vd[0],Vs[0]),min(Vd[1],Vs[1]),
          max(Vd[2],Vs[2]),min(Vd[3],Vd[3])]
    
    return dw

def dwa_control(x,u,config,goal,ob):
    dw = calc_dynamic_window(x,config)
    u,traj = calc_final_input(x,u,dw,config,goal,ob)
    return u,traj

# x = [x,y,theta,v,w]
def motion(x,u,dt):
    x[2]+= u[1]*dt
    x[0]+= u[0]*math.cos(x[2])*dt
    x[1]+= u[0]*math.sin(x[2])*dt
    x[3] = u[0]
    x[4] = u[1]
    return x



def main(gx =1 ,gy=1):
    '''
    x 当前位置，速度
    goal 目标位置
    ob 障碍物
    u 当前速度
    config机器人的基本信息
    traj 轨迹
    '''
    x = np.array([0.0, 0.0, 0.0,0.0,0.0])
    goal =np.array([gx,gy])
    ob =np.array([[-1,-1]])
    u = np.array([0.0,0.0])
    config = Config()
    traj = np.array(x)
    

    while True:
        u,ptraj = dwa_control(x,u,config,goal,ob)
        x = motion(x,u,config.dt)
        traj = np.vstack((traj,x))
        dist_to_goal = math.sqrt((x[0]-goal[0])**2+(x[1]-goal[1])**2)

        if dist_to_goal <= config.robot_radius:
            print("Goal!!")
            break

    show(traj)



def show(traj):
    x = []
    y = []
    for i in traj:
        x.append(i[0])
        y.append(i[1])

    fig = plt.figure()
    ax = fig.gca()
    figure = ax.plot(x,y)
    plt.show()
if __name__ == "__main__":
    main()